#!/usr/bin/env python
"""Django's command-line utility for administrative tasks."""
import os
import sys
import torch
import torch.nn as nn #导入nn，nn.linear是一个全连接层
from transformers import BertModel,AdamW,BertTokenizer# 导入bert模型和bert的分词器
from TorchCRF import CRF
# 问题分类模型：Bert
class BertClassfication(nn.Module):
    def __init__(self):
        # 前四行就是bert，第五行是全连接层，即bert微调
        super(BertClassfication,self).__init__()
        self.model_name = 'hfl/chinese-bert-wwm'
        self.model = BertModel.from_pretrained(self.model_name)
        self.tokenizer = BertTokenizer.from_pretrained(self.model_name)
        self.fc = nn.Linear(768,15)     #768取决于BERT结构，12-layer, 768-hidden, 12-heads, 110M parameters

    def forward(self,x):#这里的输入x是一个list,也就是输入文本x
        # 这句话是对文本1.进行分词，2.进行编码 详见https://www.cnblogs.com/zjuhaohaoxuexi/p/15135466.html
        # 第一个参数x表示文本
        # 第二个参数add_special_token表示要不要加入一些特殊字符，比如标记开头的cls等
        # 第三个参数表示最长长度
        # 第四个参数pad_to_max_len表示如果不够max_len的长度，就补长到最大长度
        # 还有一个没用到的参数参数：比最长的长久减掉，第四第五个参数也就是长的减，短的补
        batch_tokenized = self.tokenizer.batch_encode_plus(x, add_special_tokens=True,
                                max_length=148, pad_to_max_length=True)      #tokenize、add special token、pad

        # 取出input_ids：这是对汉语的编码
        # attention_mask:这是对每个字是否mask的一个标记，原本的词的位置是1，如果由于词长不够max_len，用pad或者其他的填充了，该位置就是0，表示在模型中不注意他，因为该位置是填充上的没有意义的字符
        input_ids = torch.tensor(batch_tokenized['input_ids'])
        attention_mask = torch.tensor(batch_tokenized['attention_mask'])
        # 把上边两个输入bert模型，得到bert最后一层的隐藏层的输出，
        hiden_outputs = self.model(input_ids,attention_mask=attention_mask)
        # bert的输出结果有四个维度：
        # last_hidden_state：shape是(batch_size, sequence_length, hidden_size)，hidden_size=768,它是模型最后一层输出的隐藏状态。
        # pooler_output：shape是(batch_size, hidden_size)，这是序列的第一个token(classification token)的最后一层的隐藏状态，它是由线性层和Tanh激活函数进一步处理的。（通常用于句子分类，至于是使用这个表示，还是使用整个输入序列的隐藏状态序列的平均化或池化，视情况而定）
        # hidden_states：这是输出的一个可选项，如果输出，需要指定config.output_hidden_states=True,它也是一个元组，它的第一个元素是embedding，其余元素是各层的输出，每个元素的形状是(batch_size, sequence_length, hidden_size)
        # attentions：这也是输出的一个可选项，如果输出，需要指定config.output_attentions=True,它也是一个元组，它的元素是每一层的注意力权重，用于计算self-attention heads的加权平均值。
        # 我们是微调模式，需要获取bert最后一个隐藏层的输出输入到下一个全连接层，所以取第一个维度，也就是hiden_outputs[0]
        # 此时shape是(batch_size, sequence_length, hidden_size)，[:,0,:]的意思是取出第一个也就是cls对应的结果，至于为什么这样操作，我也不知道，有人告诉我因为它最具有代表性，但为什么我还是不知道，有大神能给我讲一下吗
        outputs = hiden_outputs[0][:,0,:]     #[0]表示输出结果部分，[:,0,:]表示[CLS]对应的结果
        # 把bert最后一层的结果输入到全连接层中，全连接层是{768,15},会输出15分类的一个结果，
        output = self.fc(outputs)
        # 这里就是返回最终地分类结果了
        return output
# 命名实体识别模型：
class Model(nn.Module):
    def __init__(self,tag_num,max_length):
        super().__init__()
        self.bert = BertModel.from_pretrained('bert-base-chinese')
        self.tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
        config = self.bert.config
        self.lstm = nn.LSTM(bidirectional=True, num_layers=2, input_size=config.hidden_size, hidden_size=config.hidden_size//2, batch_first=True)
        self.crf = CRF(tag_num)
        self.fc = nn.Linear(config.hidden_size,tag_num)

    def forward(self,x):
        tokenizer_result =  self.tokenizer.encode_plus(x, return_token_type_ids=True, return_attention_mask=True,
                                                 return_tensors='pt',
                                                 padding='max_length', max_length=148)
        input_ids = torch.tensor(tokenizer_result['input_ids'])
        attention_mask = torch.tensor(tokenizer_result['attention_mask'])
        token_type_ids=torch.tensor(tokenizer_result['token_type_ids'])
        with torch.no_grad():
            bert_output = self.bert(input_ids=input_ids,attention_mask=attention_mask,token_type_ids=token_type_ids)[0]
        lstm_output, _ = self.lstm(bert_output)
        fc_output = self.fc(lstm_output)
        tag = self.crf.decode(fc_output)
        return tag

def main():
    """Run administrative tasks."""
    os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'demo.settings')
    try:
        from django.core.management import execute_from_command_line
    except ImportError as exc:
        raise ImportError(
            "Couldn't import Django. Are you sure it's installed and "
            "available on your PYTHONPATH environment variable? Did you "
            "forget to activate a virtual environment?"
        ) from exc
    execute_from_command_line(sys.argv)

if __name__ == '__main__':
    main()
